from openai import OpenAI

from Agent01.tools import tools, Tool


class SimpleAgent:
    def __init__(self):
        # 1. 定义大模型
        self.LLM = client = OpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
        # 新加坡和北京地域的API Key不同。获取API Key：https://www.alibabacloud.com/help/zh/model-studio/get-api-key
        api_key = "sk-c4b74c70f0314100a80207d8767a887e",
        # 以下为新加坡地域base_url，若使用北京地域的模型，需将base_url替换为：https://dashscope.aliyuncs.com/compatible-mode/v1
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        # 2. 定义工具
        # self.tools = {
        #     "weather": Tool(name="weather", func=getWeather, description="查询天气"),
        #     "baike": Tool(name="baike", func=WikipediaText, description="查询百科")
        # }
        self.tools = {tool.name : tool for tool in tools}

    def classify_intent(self, user_input: str) -> str:
        # 使用LLM分析用户意图
        try:
            # 定义大模型的提示词模板(prompt template)
            prompt = f"""
            分析用户输入，判断意图类型：
            - 如果是查询天气、气温、气象相关，返回"weather"
            - 如果是查询知识、定义、介绍、是什么，返回"baike" 
            - 其他情况返回"Unknown"
            只返回类型名称，不要解释。
            用户输入：{user_input}
            """
            # 调用大模型（通义千问-push）
            response = self.LLM.chat.completions.create(
                model="qwen-plus",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": f"{prompt}"},
                ],
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"意图识别错误: {e}")
            return "Unknown"

    def extract_parameter(self, user_input: str, tool_type: str) -> str:
        # 利用LLM提取执行工具需要的重要参数
        try:
            if tool_type == "weather":
                prompt = f"从文本中提取城市名称，只返回城市名：{user_input}"
            elif tool_type == "baike": # baike
                prompt = f"从文本中提取查询关键词，只返回关键词：{user_input}"
            else:
                return "未定义的工具，无法作答。"
            # 调用大模型（通义千问-push）
            response = self.LLM.chat.completions.create(
                model="qwen-plus",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": f"{prompt}"},
                ],
            )
            return response.choices[0].message.content.strip()

        except Exception as e:
            print(f"参数提取错误: {e}")
            return user_input

    def execute_tool(self, tool_name: str, parameter: str) -> str:
        # 执行工具函数
        if tool_name in self.tools:
            tool = self.tools[tool_name]
            # 动态调用工具函数, 如果tool_name是weather, 等同于执行: getWeather(parameter)
            return tool.func(parameter)
        return "未找到对应的工具，无法作答。"

    def polish_response(self, user_input: str, raw_response: str) -> str:
        # 利用LLM美化回复
        try:
            prompt = f"""
            用户问题：{user_input}
            原始回答：{raw_response}

            请将原始回答改写成更自然、友好的对话形式。
            如果原始回答比较差，请你不做作答，直接提示：智能体无法作答，请你提问天气和维基百科知识。
            """
            # 调用大模型（通义千问-push）
            response = self.LLM.chat.completions.create(
                model="qwen-plus",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": f"{prompt}"},
                ],
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"回复美化错误: {e}")
            return raw_response

    def process(self, user_input: str) -> str:
        """处理用户输入的全流程"""
        print(f"👤 用户输入: {user_input}")

        # 1. 解析意图
        intent = self.classify_intent(user_input)
        print(f"🔍 识别意图: {intent}")

        # if intent == "Unknown":
        #     return "我可以帮您查询天气或百科知识，请告诉我您想了解什么？"

        # 2. 提取参数
        parameter = self.extract_parameter(user_input, intent)
        print(f"📝 提取参数: {parameter}")

        # 3. 执行工具
        raw_result = self.execute_tool(intent, parameter)
        print(f"📊 原始结果: {raw_result}")

        # 4. 美化回复
        final_response = self.polish_response(user_input, raw_result)

        return final_response